#!/usr/bin/python3
# -*- coding: utf-8 -*-
"""
@author:  runyuanye
@contact: runyuanye@163.com
"""

import argparse
import os
import sys
import numpy as np

sys.path.append('.')

"""
   对merge的人脸ID进行过滤
"""

MIN_PITCH1 = -18.0
MAX_PITCH1 = 16.0
MIN_YAW1 = -16.0
MAX_YAW1 = 16.0
MIN_ROLL1 = -16.0
MAX_ROLL1 = 16.0

MIN_PITCH2 = -20.0
MAX_PITCH2 = 18.0
MIN_YAW2 = -18.0
MAX_YAW2 = 18.0
MIN_ROLL2 = -18.0
MAX_ROLL2 = 18.0


def read_features(file_list, feature_file, feat_dim):
    faceid_subfaceid_file_features = {}
    feat_file = open(feature_file, 'rb')
    file_count = 0
    with open(file_list, 'r') as file:
        for line in file.readlines():
            file_count += 1
            if file_count % 10000 == 0:
                print('Read Feature Count: {}'.format(file_count))

            # if file_count == 10000:
            #     break

            # 20f10478-18a0-414c-a695-c43fb51077aa 11aa6f7c-b776-4e5b-8a16-fac29fbd5cb1 81c996ed4a8a743e046fa09025d69231.jpg 0
            lines = line.strip().split()

            faceid = lines[0]
            subfaceid = lines[1]
            file_name = lines[2]
            # mode = int(lines[3])

            faceid_subfaceid = '{} {}'.format(faceid, subfaceid)

            feature = feat_file.read(feat_dim * 4)

            file_features = faceid_subfaceid_file_features.get(faceid_subfaceid, None)
            if file_features is None:
                files = [line]
                features = [feature]
                file_names = [file_name]
                faceid_subfaceid_file_features[faceid_subfaceid] = [files, features, file_names]
            else:
                file_features[0].append(line)
                file_features[1].append(feature)
                file_features[2].append(file_name)

        print('Read Feature Count: {}'.format(file_count))
    feat_file.close()
    return faceid_subfaceid_file_features


def read_attr_file(file_list):
    file_count = 0
    name_attrs = {}
    with open(file_list, 'r') as file:
        for line in file.readlines():
            file_count += 1
            if file_count % 10000 == 0:
                print('Read Attr Count: {}'.format(file_count))

            # if file_count == 10000:
            #     break

            # gd/zs/nqjd/dxjyjy/150100414a54443452064d2742c93600/video/20190817_082812/000359_000.jpg;-12.010,20.456,50.344;0.630
            lines = line.strip().split(';')
            name = lines[0]
            headpose = lines[1].split(',')
            headpose = list(map(float, headpose))  # pitch, roll, yaw
            quality = float(lines[2])
            name_attrs[name] = [headpose, quality]
        print('Read Attr Count: {}'.format(file_count))
    return name_attrs


def main():
    parser = argparse.ArgumentParser(description="Replace Face ReID Image")
    parser.add_argument(
        "--file_list", default="/rootfs/media/kasim/Data1/data/VideoFaceImage/filter3/FaceRecSel.txt", help="path to image file list", type=str
    )
    parser.add_argument(
        "--feature_file", default="/rootfs/media/kasim/Data1/data/VideoFaceImage/filter3/FaceRecFeatSel.dat", help="path to feature file", type=str
    )
    parser.add_argument(
        "--attr_file", default="/rootfs/media/kasim/Data1/data/VideoFaceImage/filter3/FaceAttrs.txt", help="path to merge file", type=str
    )
    parser.add_argument(
        "--output_dir", default="/rootfs/media/kasim/Data1/data/VideoFaceImage/filter3", help="path to output", type=str
    )
    parser.add_argument(
        "--feature_dim", default=256, help="feature dim", type=int
    )
    args = parser.parse_args()

    output_dir = args.output_dir
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
        os.system('chmod a+wr {}'.format(output_dir))

    faceid_subfaceid_file_feature_names = read_features(args.file_list, args.feature_file, args.feature_dim)
    name_attrs = read_attr_file(args.attr_file)

    out_file_list_path = os.path.join(output_dir, 'TypicalFace.txt')
    out_feature_file_path = os.path.join(output_dir, 'TypicalFaceFeature.dat')
    out_file_list = open(out_file_list_path, 'w')
    out_feature_file = open(out_feature_file_path, 'wb')
    face_count = 0
    face_image_count = 0
    for faceid_subfaceid, file_feature_names in faceid_subfaceid_file_feature_names.items():
        file_infos, features, file_names = file_feature_names
        attrs = [name_attrs[name] for name in file_names]
        pitchs = []
        rolls = []
        yaws = []
        qualitys = []
        for [pitch, roll, yaw], quality in attrs:
            pitchs.append(pitch)
            rolls.append(roll)
            yaws.append(yaw)
            qualitys.append(quality)
        pitchs = np.array(pitchs, dtype=np.float32)
        rolls = np.array(rolls, dtype=np.float32)
        yaws = np.array(yaws, dtype=np.float32)
        qualitys = np.array(qualitys, dtype=np.float32)

        image_count = len(attrs)
        face_count += 1
        face_image_count += image_count
        idxes = np.arange(0, image_count)
        pitchs_mask = np.logical_and(MIN_PITCH1 <= pitchs, pitchs <= MAX_PITCH1)
        rolls_mask = np.logical_and(MIN_ROLL1 <= rolls, rolls <= MAX_ROLL1)
        yaws_mask = np.logical_and(MIN_YAW1 <= yaws, yaws <= MAX_YAW1)
        mask = pitchs_mask & rolls_mask & yaws_mask
        if mask.any():
            max_quality_idx = qualitys[mask].argmax()
            idx = idxes[mask][max_quality_idx]
        else:
            pitchs_mask = np.logical_and(MIN_PITCH2 <= pitchs, pitchs <= MAX_PITCH2)
            rolls_mask = np.logical_and(MIN_ROLL2 <= rolls, rolls <= MAX_ROLL2)
            yaws_mask = np.logical_and(MIN_YAW2 <= yaws, yaws <= MAX_YAW2)
            mask = pitchs_mask & rolls_mask & yaws_mask
            if mask.any():
                max_quality_idx = qualitys[mask].argmax()
                idx = idxes[mask][max_quality_idx]
            else:
                idx = qualitys.argmax()

        file_info = '{} {} {}\n'.format(faceid_subfaceid, file_names[idx], image_count)
        feature = features[idx]
        out_file_list.write(file_info)
        out_feature_file.write(feature)
    out_feature_file.close()
    out_file_list.close()
    os.system('chmod a+wr {}'.format(out_file_list_path))
    os.system('chmod a+wr {}'.format(out_feature_file_path))
    print('Face Count: {}, Face Image Count: {}'.format(face_count, face_image_count))
    print('finish!')


if __name__ == '__main__':
    main()

